import cv2
import numpy as np


# 初始化背景减法器
background_subtractor = cv2.createBackgroundSubtractorMOG2(history=500, varThreshold=16, detectShadows=False)

# 打开视频文件
# cap = cv2.VideoCapture('./minicar.mp4')  # 这里需要替换成你的视频路径
cap = cv2.VideoCapture('./rc.mp4')  # 这里需要替换成你的视频路径

# 初始化变量
points = []
frame_index = 0

# 检查视频是否打开成功
if not cap.isOpened():
    print("Error: Could not open video.")
    exit()

# 定义白色颜色范围（在 HSV 颜色空间中）
lower_white = np.array([0, 0, 200])  # HSV下限
upper_white = np.array([180, 30, 255])  # HSV上限


while True:
    ret, frame = cap.read()
    if not ret:
        # 如果视频短小，播放完毕把视频指针重置到初始的位置，就是特喵的循环播放
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        continue

    # 转换为灰度图像
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)

    # 应用背景减法
    fg_mask = background_subtractor.apply(gray)

    # 应用阈值来获取二值图像
    _, thresh = cv2.threshold(fg_mask, 25, 255, cv2.THRESH_BINARY)

    # 添加形态学操作以减少噪声
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)

    # 转换到 HSV 颜色空间，用于白色检测
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    white_mask = cv2.inRange(hsv, lower_white, upper_white)

    # 将白色掩码与前景掩码结合
    combined_mask = cv2.bitwise_and(thresh, white_mask)

    # 查找轮廓
    contours, _ = cv2.findContours(combined_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 遍历轮廓
    for contour in contours:
        # 过滤掉小轮廓
        if cv2.contourArea(contour) < 20 or cv2.contourArea(contour) > 100:
            continue

        # 计算边界框
        x, y, w, h = cv2.boundingRect(contour)

        # 绘制边界框
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)

        # 计算中心点
        center_x = x + w // 2
        center_y = y + h // 2

        # 保存中心点用于绘制轨迹
        points.append((center_x, center_y))

    # 绘制轨迹
    for i in range(1, len(points)):
        # 如果两个点之间的距离过大，跳过
        if abs(points[i][0] - points[i - 1][0]) < 50 and abs(points[i][1] - points[i - 1][1]) < 50:
            # 用红色绘制轨迹
            cv2.line(frame, points[i - 1], points[i], (0, 0, 255), 2)

    # 显示结果
    cv2.imshow('Frame', frame)

    # 按 'q' 键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

    frame_index += 2

# 释放资源
cap.release()
cv2.destroyAllWindows()
